#! /usr/bin/env python
# coding=utf-8
seq_file_name_aa = 'pep_cut_ali.aa'



import os
from Bio import SeqIO
import argparse
import sys

# 使用方法 
# conver04_add_gap2simulation.py  -r rawfile -s simulationfile 

parser = argparse.ArgumentParser(
    description=''' 将gap 加入模拟
    用法:
    conver04_add_gap2simulation.py  -r OGS_for_second_step -s simulate0
    由大天才于2021年12月31日创建于浙江农业大学''')

parser.add_argument('-r',
                help='用于测试的 ogs文件路径')

parser.add_argument('-s',
                help='模拟的 ogs文件路径')

args = parser.parse_args()


if not args.r or not args.s:
    parser.print_help()
    sys.exit()






rawfile = args.r

simulationfile = args.s




# rawfile = r'OGS_for_second_step'
# simulationfile = 'simulate1'

n = 0

for ogs_id in os.listdir(simulationfile):
	n+=1


	if not os.path.exists(simulationfile+'/'+ogs_id+'/'+'aaml.ctl'):
		continue
	else:
		if n%50==0:
			print('.')



	# ogs_id = 'OG0002825'
	raw_path = rawfile+'/'+ogs_id+'/'+'pep_cut_ali.fa'
	simulate_path = simulationfile + '/' + ogs_id+'/pep_cut_ali.fa'
	raw_seq = {str(i.name):str(i.seq) for i in SeqIO.parse(raw_path,'fasta')}

	simulate_seq = {str(i.name):str(i.seq) for i in SeqIO.parse(simulate_path,'fasta')}
	withgap_file = open(simulationfile + '/' + ogs_id+'/pep_cut_ali_with_gap.fa','w')
	# 重新写入  pep_cut_ali.aa

	pep_cut_ali_file = open(simulationfile + '/' + ogs_id+'/'+seq_file_name_aa,'w')
	for i in simulate_seq:
		break
	seq_nub = str(len(raw_seq))
	
	seq_len = str(len(simulate_seq[i]))

	pep_cut_ali_file.write(seq_nub+'  '+seq_len+'\n')
	pep_cut_ali_file.write('\n')

	for i in raw_seq:
		new_seq = simulate_seq[i]

		for j in range(len(raw_seq[i])):
			if raw_seq[i][j] =='-':
				new_seq = new_seq[:j]+ '-'+new_seq[j+1:]
		withgap_file.write('>'+i+'\n'+new_seq+'\n')
		pep_cut_ali_file.write('>\n'+i+'\n'+new_seq+'\n\n')
		#pep_cut_ali_file.write()


	withgap_file.close()
	pep_cut_ali_file.close()

	# break

		#simulate_seq

	#break